#!/usr/bin/python
# -*- encoding: utf-8 -*-

"""
@File:      tc_ai_pytorch_quicik_start.py
@Time:      2023/07/4 11:12:50
@Author:    zhangtaibo
@Version:   1.0
@Contact:   zhangtaibo.ztb@alibaba-inc.com
@License:   Mulan PSL v2
@Modify:    zhangtaibo
"""
from common.basetest import LocalTest
class Test(LocalTest):
    """
    See tc_ai_pytorch_quici_start.yaml for details

    :avocado: tags=fix,P1,noarch,local,pytorch_ai_container
    """

    def setUp(self):
        super().setUp()
        ret_c,self.docker_name = self.cmd(self.container_engine+" ps |grep %s |awk '{print $NF}'" % self.version)
        self.cmd(self.container_engine+" cp res/ai/pytorch/create_csv.py %s:/tmp" % self.docker_name)
        self.cmd(self.container_engine+" cp res/ai/pytorch/pytorch_quit_start_csv.py %s:/tmp" % self.docker_name)
        self.setup_rpm_install("git", flag=1)
        self.cmd("pip3 install pandas",container_flag=1)
        self.cmd("mkdir -p ./data",container_flag=1)

    def test(self):
        self.cmd("python3 /tmp/create_csv.py",container_flag=1)
        self.cmd("ls ./data/train.csv",container_flag=1)
        _, output = self.cmd("python3 /tmp/pytorch_quit_start_csv.py",container_flag=1)
        self.cmd("ls model.ckpt",container_flag=1)
        self.assertTrue("Epoch" in output, 'pytorch quit start test failed, please check.')

    def tearDown(self):
        super().tearDown()

